#!/usr/bin/python3
# -*- coding: utf-8; tab-width: 4; indent-tabs-mode: t -*-

import os
import glob
import subprocess
import strict_systemd
from fm_util import Util


class SwapManager:

    def __init__(self, layout, bSimple):
        self._layout = layout
        self._bSimple = bSimple
        self._cfgDirObj = strict_systemd.cfg.SysCfgDir()

    def getInfo(self):
        if self._layout.name in ["bios-ext4", "efi-ext4"]:
            return self._getInfoOnDisk()
        else:
            return self._getInfoInMemory()

    def check(self, bAutoFix, errorCallback):
        # check swappiness
        if Util.readFlagFile("/proc/sys/vm/swappiness") not in ["0", "1"]:
            # don't abort, we may still be able to use swap
            errorCallback("Swappiness should be 0 or 1.")

        if self._layout.name in ["bios-ext4", "efi-ext4"]:
            if not self._checkSwapOnDisk(bAutoFix, errorCallback):
                return
            swapServiceCount = 1
        else:
            if not self._checkSwapInMemory(bAutoFix, errorCallback):
                return
            swapServiceCount = 0

        if len(self._cfgDirObj.query_services(filter=lambda fn, fullfn: fn.endswith(".swap"))) != swapServiceCount:
            errorCallback("Redundant swap service found.")

    def _getInfoOnDisk(self):
        if self._layout.has_swap_file():
            try:
                swapObj = strict_systemd.runtime.SwapService(True, swap_path=self._layout.get_swap_file_path())
                if swapObj.is_running():
                    return (True, os.path.getsize(self._layout.get_swap_file_path()))
            except FileNotFoundError:
                return None
        else:
            return None

    def _getInfoInMemory(self):
        try:
            swapObj = strict_systemd.runtime.DaemonService(True, name="zram-swap")
            if swapObj.is_running():
                return (False, Util.zramDevGetSize("/dev/zram0"))
        except FileNotFoundError:
            return None

    def _checkSwapOnDisk(self, bAutoFix, errorCallback):
        # our swap space is limited, we have to use zswap
        if not os.path.exists("/sys/module/zswap"):
            errorCallback("zswap is not enabled.")
            return False

        # our swap space is limited, we have to use zstd compressor
        if Util.readFlagFile("/sys/module/zswap/parameters/compressor") != "zstd":
            errorCallback("zswap should use zstd as compressor.")
            return False

        # check swap file existence
        if not self._layout.has_swap_file():
            if bAutoFix:
                self.__enableSwapOnDisk()
            else:
                errorCallback("Swap file not found.")
                return False

        swapSize = os.path.getsize(self._layout.get_swap_file_path())

        # check swap file size
        if swapSize < Util.getOnDiskSwapSizeInGb() * 1024 * 1024 * 1024:
            if bAutoFix:
                self.__disableSwapOnDisk()            # disable and enable so that we re-create a new swap file
                self.__enableSwapOnDisk()
            else:
                # don't abort, we may still be able to use this swap file
                errorCallback("Swap file is too small.")

        # check swap file size
        if swapSize < int(Util.readFlagFile("/sys/power/image_size", "0")):
            if bAutoFix:
                self.__disableSwapOnDisk()        # disable and enable so that we re-create a new swap file
                self.__enableSwapOnDisk()
            else:
                # don't abort, we may still be able to use this swap file
                errorCallback("Swap file size is smaller than /sys/power/image_size.")

        # check zswap enablement
        if not Util.strYnToBool(Util.readFlagFile("/sys/module/zswap/parameters/enabled")):
            if bAutoFix:
                self.__enableSwapOnDisk()
            else:
                errorCallback("zswap is not enabled.")
                return False

        # swap service should exist
        swapCfgObj = self._cfgDirObj.try_get_service(strict_systemd.ServiceType.SWAP, swap_path=self._layout.get_swap_file_path())
        if swapCfgObj is None:
            if bAutoFix:
                swapObj = self.__enableSwapOnDisk()
            else:
                errorCallback("Swap service not found.")
                return False
        else:
            swapObj = strict_systemd.runtime.SwapService(True, cfg=swapCfgObj)

        # swap service should in /etc
        if not swapObj.unit_file_path.startswith(self._cfgDirObj.etc_system_dir + "/"):
            # don't abort, we may still be able to use this swap file
            errorCallback("Swap service should be located in /etc.")

        # swap service should be enabled
        if not swapObj.is_enabled():
            if bAutoFix:
                self.__enableSwapOnDisk()
            else:
                errorCallback("Swap serivce is not enabled.")
                return False

        # swap service should be started
        if not self._bSimple:
            if not swapObj.is_running():
                if bAutoFix:
                    self.__enableSwapOnDisk()
                else:
                    errorCallback("Swap service is not started.")
                    return False
        else:
            lineList = subprocess.check_output(["swapon", "-s"], text=True).split("\n")
            if len(lineList) < 3:
                if bAutoFix:
                    self.__enableSwapOnDisk()
                else:
                    errorCallback("Swap service is not started.")
                    return False
            elif len(lineList) == 3:
                if lineList[1].split()[0] != self._layout.get_swap_file_path():
                    errorCallback("The running swap is invalid.")
                    return False
            else:
                errorCallback("The running swap is invalid.")
                return False

        return True

    def _checkSwapInMemory(self, bAutoFix, errorCallback):
        # zswap must not be enabled
        if Util.readFlagFile("/sys/module/zswap/parameters/enabled", "N") != "N":
            errorCallback("zswap should not be enabled.")
            return False

        # check redundant zram devices
        if glob.glob("/dev/zram*") not in [[], ["/dev/zram0"]]:
            errorCallback("Multiple zram devices exist.")
            return False

        # check swap file existence (symlink exists and symlink points to a valid file)
        if not os.path.exists("/dev/zram0"):
            if bAutoFix:
                self.__enableSwapInMemory()
            else:
                errorCallback('"/dev/zram0" not found.')
                return False

        # check swap file size
        if Util.zramDevGetSize("/dev/zram0") < Util.getInMemorySwapSizeInMb() * 1024 * 1024:
            if bAutoFix:
                self.__disableSwapInMemory()          # disable and enable so that we re-create a new zram device
                self.__enableSwapInMemory()
            else:
                # don't abort, we may still be able to use this swap file
                errorCallback('"/dev/zram0" has too small zram disk size.')

        # swap service should exist
        servCfgObj = self._cfgDirObj.try_get_service(strict_systemd.ServiceType.DAEMON, name="zram-swap")
        if servCfgObj is None:
            if bAutoFix:
                servObj = self.__enableSwapInMemory()
            else:
                errorCallback("Swap service not found.")
                return False
        else:
            servObj = strict_systemd.runtime.DaemonService(True, cfg=servCfgObj)

        # swap service should in /etc
        if not servObj.unit_file_path.startswith(self._cfgDirObj.etc_system_dir + "/"):
            # don't abort, we may still be able to use this swap file
            errorCallback("Swap service should be located in /etc.")

        # swap service should be enabled
        if not servObj.is_enabled():
            if bAutoFix:
                self.__enableSwapInMemory()
            else:
                errorCallback("Swap serivce is not enabled.")
                return False

        # swap service should be started
        if not self._bSimple:
            if not servObj.is_running():
                if bAutoFix:
                    self.__enableSwapInMemory()
                else:
                    errorCallback("Swap service is not started.")
                    return False
        else:
            lineList = subprocess.check_output(["swapon", "-s"], text=True).split()
            if len(lineList) < 3:
                if bAutoFix:
                    self.__enableSwapInMemory()
                else:
                    errorCallback("Swap service is not started.")
                    return False
            elif len(lineList) == 3:
                if lineList[1].split()[0] != "/dev/zram0":
                    errorCallback("The running swap is invalid.")
                    return False
            else:
                errorCallback("The running swap is invalid.")
                return False

        return True

    def __enableSwapOnDisk(self):
        if not self._layout.has_swap_file():
            self._layout.create_swap_file("%dG" % (Util.getOnDiskSwapSizeInGb()))

        Util.writeFlagFile("/sys/module/zswap/parameters/enabled", "Y")

        swapCfgObj = self._cfgDirObj.create_service(strict_systemd.ServiceType.SWAP, self._cfgDirObj.etc_system_dir, exist_ok=True, swap_path=self._layout.get_swap_file_path())
        swapObj = strict_systemd.runtime.SwapService(True, cfg=swapCfgObj)
        swapObj.enable()

        if not self._bSimple:
            swapObj.start()
        else:
            subprocess.check_call(["swapon", self._layout.get_swap_file_path()])

        return swapObj

    def __disableSwapOnDisk(self):
        swapCfgObj = self._cfgDirObj.try_get_service(strict_systemd.ServiceType.SWAP, swap_path=self._layout.get_swap_file_path())
        if swapCfgObj is not None:
            swapObj = strict_systemd.runtime.SwapService(True, cfg=swapCfgObj)
            if not self._bSimple:
                swapObj.stop()
            else:
                subprocess.check_call(["swapoff", self._layout.get_swap_file_path()])
            swapObj.disable()
            swapObj.remove_and_dispose()

        Util.writeFlagFile("/sys/module/zswap/parameters/enabled", "N")

        self._layout.remove_swap_file()

    def __enableSwapInMemory(self):
        if os.path.exists("/dev/zram0"):
            self.__disableSwapInMemory()

        buf = ""
        buf += "[Unit]\n"
        buf += "Description=Swap with zram\n"
        buf += "\n"
        buf += "[Service]\n"
        buf += "Type=oneshot\n"
        buf += "RemainAfterExit=true\n"
        buf += "ExecStart=/usr/libexec/fpemud-os-sysman/zram-swap.py init\n"
        buf += "ExecStop=/usr/libexec/fpemud-os-sysman/zram-swap.py fini\n"
        buf += "\n"
        buf += "[Install]\n"
        buf += "WantedBy=swap.target\n"

        servCfgObj = self._cfgDirObj.create_service(strict_systemd.ServiceType.DAEMON, self._cfgDirObj.etc_system_dir, exist_ok=True, name="zram-swap", content=buf)
        servObj = strict_systemd.runtime.DaemonService(True, cfg=servCfgObj)
        servObj.enable()

        if not self._bSimple:
            servObj.start()
        else:
            subprocess.check_call(["/usr/libexec/fpemud-os-sysman/zram-swap.py", "init"])

        return servObj

    def __disableSwapInMemory(self):
        servCfgObj = self._cfgDirObj.try_get_service(strict_systemd.ServiceType.DAEMON, name="zram-swap")
        if servCfgObj is not None:
            servObj = strict_systemd.runtime.DaemonService(True, cfg=servCfgObj)
            if not self._bSimple:
                servObj.stop()
            servObj.disable()
            servObj.remove_and_dispose()

        # as a "close" action, ensure we are closed
        if os.path.exists("/dev/zram0"):
            subprocess.check_call(["/usr/libexec/fpemud-os-sysman/zram-swap.py", "fini"])
